require "helper"
require "socket"
require "zlib"

class TestRackUpstreamIdentification < Test::Unit::TestCase
  include Rack::Test::Methods

  def app
    handler = lambda { |env| [200, {"Content-Type" => "text/plain"}, ["OK"]] }
    Rack::Builder.new do
      map "/default" do
        use Rack::UpstreamIdentification
        run handler
      end

      map "/override_host" do
        use Rack::UpstreamIdentification, "server_header"
        run handler
      end

      map "/transform_block" do
        use Rack::UpstreamIdentification do |host|
          host.upcase
        end
        run handler
      end

      map "/both_options" do
        use Rack::UpstreamIdentification, "crypted_header" do |host|
          Zlib.crc32(host)
        end
        run handler
      end

      map "/empty_block" do
        use Rack::UpstreamIdentification do |host|
        end
        run handler
      end

    end.to_app
  end

  def assert_header_present(header)
    assert last_response.headers.include? header
  end

  def assert_header_value(header, value)
    assert_equal value, last_response.headers[header]
  end

  context "using the upstream identification handler" do

    should "allow requests to flow through without issue" do
      get "/default"
      assert last_response.ok?
    end

    context "with the default configuration" do
      setup do
        get "/default"
      end

      should "push the header down using the default name" do
        assert_header_present "X-Upstream-Identification"
      end

      should "push the header down using the machine's hostname" do
        assert_header_value "X-Upstream-Identification", Socket.gethostname
      end

    end

    context "with an overriding the header name" do
      setup do
        get "/override_host"
      end

      should "push the header down with the overriden name" do
        assert_header_present "server_header"
      end

      should "push the header down using the machine's hostname" do
        assert_header_value "server_header", Socket.gethostname
      end

    end

    context "with a block to provide custom processing" do
      setup do
        get "/transform_block"
      end

      should "push the header down with the default header name" do
        assert_header_present "X-Upstream-Identification"
      end

      should "push the header down using the transformed value" do
        assert_header_value "X-Upstream-Identification", Socket.gethostname.upcase
      end

    end

    context "with the headername and transform block passed" do
      setup do
        get "/both_options"
      end

      should "push the header down with the overridden header name" do
        assert_header_present "crypted_header"
      end

      should "push the header down using the mutated header value" do
        assert_header_value "crypted_header", Zlib.crc32(Socket.gethostname).to_s
      end

    end

    context "with an empty transform block" do
      should "raise an ArgumentError indicating that the transform block should return a non-nil value" do
        assert_raise ArgumentError do
          get "/empty_block"
        end
      end

    end

  end

end
